import numpy as np
from scipy.sparse import coo_matrix, csc_matrix, csr_matrix, spdiags, eye, bmat


class NatureBC():
    def __init__(self, space, gD, threshold=None):
        self.space = space
        self.gD = gD
        self.threshold = threshold
        self.bctype = 'NatureBC'
    def apply(self,F,threshold=None,q=None):
        """
        Parameters
        ----------

        """            
        space = self.space
        gD = self.gD
        threshold = self.threshold if threshold is None else threshold
        space.set_nature_bc(F,gD,threshold=threshold,q=q)
        return F


class EssentialBC():
    '''
    本质边界条件
    '''
    def __init__(self,space,gN,threshold=None):
        self.space = space
        self.gN = gN
        self.threshold = threshold
        self.bctype = 'Neumann'

    def apply(self, A, F, uh=None, threshold=None):
        space = self.space
        gN = self.gN
        threshold = self.threshold if threshold is None else threshold

        if uh is None:
            uh = self.space.function()
        isBdDof = space.set_essential_bc(uh,gN,A=A,b=F,threshold=threshold) #（tensorgdof,）
        self.isBdDof = isBdDof
        #print(np.max(np.abs(A-A.T)),F)
        gdof = list(A.shape)[0]
        tensorgdof = list(uh.shape)[0]
        x = np.zeros(gdof,dtype=float)
        x[:tensorgdof][isBdDof] = uh[isBdDof]
        x = x.T.flat
        F -= A@x
        bdIdx = np.zeros(gdof, dtype=np.int)
        bdIdx[:tensorgdof][isBdDof] = 1
        Tbd = spdiags(bdIdx, 0, A.shape[0], A.shape[0])
        T = spdiags(1-bdIdx, 0, A.shape[0], A.shape[0])  
        A = T@A@T + Tbd
        #print(np.max(np.abs(F)))
        F[:tensorgdof][isBdDof] = x[:tensorgdof][isBdDof]
        #print(np.max(np.abs(F)))
        return A, F         
